#include <iostream>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <sys/time.h>
#include "mkl.h"
#include <omp.h>
using namespace std;
double ***A;
double * Alin;

#define real double
#define integer int
int inline dgemm2(char *transa, char *transb, integer *m, integer *n, integer *k, real *alpha,
                        real *a, integer *lda, real *b, integer *ldb, real *beta, real *c__, integer *ldc)
{
    integer a_dim1, a_offset, b_dim1, b_offset, c_dim1, c_offset, i__1, i__2,i__3;
    real temp;
    integer i__, j, l;

    a_dim1 = *lda;
    a_offset = 1 + a_dim1 * 1;
    a -= a_offset;
    b_dim1 = *ldb;
    b_offset = 1 + b_dim1 * 1;
    b -= b_offset;
    c_dim1 = *ldc;
    c_offset = 1 + c_dim1 * 1;
    c__ -= c_offset;

    i__1 = *n;
    for (j = 1; j <= i__1; ++j)
    {
        i__2 = *k;
        for (l = 1; l <= i__2; ++l)
        {
            temp = *alpha * b[j + l * b_dim1];
            i__3 = *m;
            for (i__ = 1; i__ <= i__3; ++i__)
                c__[i__ + j * c_dim1] += temp * a[i__ + l * a_dim1];
        }
    }
    return 0;
}

int inline dtrsm2(char *side, char *uplo, char *transa, char *diag, integer *m,
                        integer *n, real *alpha, real *a, integer *lda, real *b, integer *ldb)
{
    integer a_dim1, a_offset, b_dim1, b_offset, i__1, i__2, i__3;
    real temp;
    integer i__, j, k;

    a_dim1 = *lda;
    a_offset = 1 + a_dim1 * 1;
    a -= a_offset;
    b_dim1 = *ldb;
    b_offset = 1 + b_dim1 * 1;
    b -= b_offset;

    i__1 = *n;
    for (k = 1; k <= i__1; ++k)
    {
        temp = 1.f / a[k + k * a_dim1];
        i__2 = *m;
        for (i__ = 1; i__ <= i__2; ++i__)
            b[i__ + k * b_dim1] = temp * b[i__ + k * b_dim1];
        i__2 = *n;
        for (j = k + 1; j <= i__2; ++j)
        {
            temp = a[j + k * a_dim1];
            i__3 = *m;
            for (i__ = 1; i__ <= i__3; ++i__)
                b[i__ + j * b_dim1] -= temp * b[i__ + k * b_dim1];
        }
    }
    return 0;
}

int inline dsyrk2(char *uplo, char *trans, integer *n, integer *k, real *alpha,
                        real *a, integer *lda, real *beta, real *c__, integer *ldc)
{
    integer a_dim1, a_offset, c_dim1, c_offset, i__1, i__2, i__3;
    real temp;
    integer i__, j, l;

    a_dim1 = *lda;
    a_offset = 1 + a_dim1 * 1;
    a -= a_offset;
    c_dim1 = *ldc;
    c_offset = 1 + c_dim1 * 1;
    c__ -= c_offset;

    i__1 = *n;
    for (j = 1; j <= i__1; ++j)
    {
        i__2 = *k;
        for (l = 1; l <= i__2; ++l)
        {
            temp = *alpha * a[j + l * a_dim1];
            i__3 = *n;
            for (i__ = j; i__ <= i__3; ++i__)
                c__[i__ + j * c_dim1] += temp * a[i__ + l * a_dim1];
        }
    }
    return 0;
}
#undef real
#undef integer

// -------------------------------------------------------------------------

void dpo(double *T, int N)
{
    int i, j, k;
    for(k = 0; k < N; k++)
    {
        double sum = 0;
        for(i = 0; i < k; i++)
            sum += T[i*N+k] * T[i*N+k];
        sum = T[k*N+k] - sum;
        T[k*N+k] = sqrt(sum > 0 ? sum : 0);
        for(i = k + 1; i < N; i++)
        {
            sum = 0;
            for(j = 0; j < k; j++)
                sum += T[j*N+i] * T[j*N+k];
            T[k*N+i] = (T[i*N+k] - sum) / T[k*N+k];
        }
    }

}


#define real double
#define integer int
int myspotf2_(char *uplo, integer *n, real *a, integer *lda, integer *info)
{
    static integer c__1 = 1;
    static real c_b10 = -1.f;
    static real c_b12 = 1.f;
    integer a_dim1, a_offset, i__1, i__2, i__3;
    real r__1;

    static integer j;
    static real ajj;
    extern double mysdot_(integer *, real *, integer *, real *, integer *);
    extern int mysscal_(integer *, real *, real *, integer *),
           mysgemv_(char *, integer *, integer *, real *, real *, integer *,
                    real *, integer *, real *, real *, integer *);
    a_dim1 = *lda;
    a_offset = 1 + a_dim1;
    a -= a_offset;
    i__1 = *n;
    for (j = 1; j <= i__1; ++j)
    {
        i__2 = j - 1;
        ajj = a[j + j * a_dim1] - mysdot_(&i__2, &a[j + a_dim1], lda, &a[j + a_dim1], lda);
        if (ajj <= 0.f)
        {
            a[j + j * a_dim1] = ajj;
            *info = j;
            return 0;
        }
        ajj = sqrt(ajj);
        a[j + j * a_dim1] = ajj;
        if (j < *n)
        {
            i__2 = *n - j;
            i__3 = j - 1;
            mysgemv_("No transpose", &i__2, &i__3, &c_b10, &a[j + 1 + a_dim1], lda, &a[j + a_dim1],
                     lda, &c_b12, &a[j + 1 + j * a_dim1], &c__1);
            i__2 = *n - j;
            r__1 = 1.f / ajj;
            mysscal_(&i__2, &r__1, &a[j + 1 + j * a_dim1], &c__1);
        }
    }
    return 0;
}

double mysdot_(integer *n, real *sx, integer *incx, real *sy, integer *incy)
{
    integer i__1;
    real ret_val;
    static integer i__, m, ix, iy, mp1;
    static real stemp;
    --sy;
    --sx;
    stemp = 0.f;
    ret_val = 0.f;
    if (*n <= 0)
    {
        return ret_val;
    }
    if (*incx == 1 && *incy == 1)
    {
        goto L20;
    }

    ix = 1;
    iy = 1;

    i__1 = *n;
    for (i__ = 1; i__ <= i__1; ++i__)
    {
        stemp += sx[ix] * sy[iy];
        ix += *incx;
        iy += *incy;
    }
    ret_val = stemp;
    return ret_val;

L20:
    m = *n % 5;
    if (m == 0)
    {
        goto L40;
    }
    i__1 = m;
    for (i__ = 1; i__ <= i__1; ++i__)
    {
        stemp += sx[i__] * sy[i__];
    }
    if (*n < 5)
    {
        ret_val = stemp;
        return ret_val;
    }

L40:
    mp1 = m + 1;
    i__1 = *n;
    for (i__ = mp1; i__ <= i__1; i__ += 5)
    {
        stemp = stemp + sx[i__] * sy[i__] + sx[i__ + 1] * sy[i__ + 1] + sx[i__ + 2] * sy[i__ + 2]
                + sx[i__ + 3] * sy[i__ + 3] + sx[i__ + 4] * sy[i__ + 4];
    }
}
int mysscal_(integer *n, real *sa, real *sx, integer *incx)
{
    integer i__1, i__2;
    static integer i__, m, mp1, nincx;
    --sx;
    if (*n <= 0 || *incx <= 0)
    {
        return 0;
    }
    if (*incx == 1)
    {
        goto L20;
    }
    nincx = *n **incx;
    i__1 = nincx;
    i__2 = *incx;
    for (i__ = 1; i__2 < 0 ? i__ >= i__1 : i__ <= i__1; i__ += i__2)
    {
        sx[i__] = *sa * sx[i__];
    }
    return 0;

L20:
    m = *n % 5;
    if (m == 0)
    {
        goto L40;
    }
    i__2 = m;
    for (i__ = 1; i__ <= i__2; ++i__)
    {
        sx[i__] = *sa * sx[i__];
    }
    if (*n < 5)
    {
        return 0;
    }

L40:
    mp1 = m + 1;
    i__2 = *n;
    for (i__ = mp1; i__ <= i__2; i__ += 5)
    {
        sx[i__] = *sa * sx[i__];
        sx[i__ + 1] = *sa * sx[i__ + 1];
        sx[i__ + 2] = *sa * sx[i__ + 2];
        sx[i__ + 3] = *sa * sx[i__ + 3];
        sx[i__ + 4] = *sa * sx[i__ + 4];
    }
    return 0;
}

int mysgemv_(char *trans, integer *m, integer *n, real *alpha, real *a,
             integer *lda, real *x, integer *incx, real *beta, real *y, integer *incy)
{
    integer a_dim1, a_offset, i__1, i__2;
    static integer i__, j, jx;
    static real temp;
    a_dim1 = *lda;
    a_offset = 1 + a_dim1;
    a -= a_offset;
    --x;
    --y;
    jx = 1;

    i__1 = *n;
    for (j = 1; j <= i__1; ++j)
    {
        if (x[jx] != 0.f)
        {
            temp = *alpha * x[jx];
            i__2 = *m;
            for (i__ = 1; i__ <= i__2; ++i__)
            {
                y[i__] += temp * a[i__ + j * a_dim1];
            }
        }
        jx += *incx;
    }
    return 0;
}
#undef real
#undef integer


#pragma oss task inout([NB*NB]T)
void dpotrf_tile(double *T, int NB)
{
    int info;
    myspotf2_("L", &NB, T, &NB, &info);
}

#pragma oss task in([NB*NB]T, [NB*NB]B) inout([NB*NB]C)
void dgemm_tile(double *T, double *B, double *C, int NB)
{
    double done=1.0, dmone=-1.0;
    char TR='T', NT='N';
    dgemm2(&NT, &TR, &NB, &NB, &NB, &dmone, T, &NB, B, &NB, &done, C, &NB);
}

#pragma oss task in([NB*NB]T) inout([NB*NB]B)
void dtrsm_tile(double *T, double *B, int NB)
{
    double done=1.0;
    char LO='L', TR='T', NU='N', RI='R';
    dtrsm2(&RI, &LO, &TR, &NU, &NB, &NB, &done, T, &NB, B, &NB);
}

#pragma oss task in([NB*NB]T) inout([NB*NB]B)
void dsyrk_tile( double *T, double *B, int NB)
{
    double done=1.0, dmone=-1.0;
    char LO='L', NT='N';
    dsyrk2(&LO, &NT, &NB, &NB, &dmone, T, &NB, &done, B, &NB);
}

void compute(int NB, int DIM, double ***A)
{
    int j,k,i;
    int N = NB*NB;
    for (j = 0; j < DIM; j++)
    {
        for (k= 0; k< j; k++)
        {
            for (i = j+1; i < DIM; i++)
            {
                dgemm_tile( A[k][i], A[k][j], A[j][i], NB);
            }
        }

        for (i = 0; i < j; i++)
        {
            dsyrk_tile( A[i][j], A[j][j], NB);
        }

        dpotrf_tile( A[j][j], NB);

        for (i = j+1; i < DIM; i++)
        {
            dtrsm_tile( A[j][j], A[j][i], NB);
        }
    }
    #pragma oss taskwait
}

//--------------------------------------------------------------------------------

int check(int NB, int N, double *Alin, double ***A)
{
    int i,j;
    for (i = 0; i < N; i++)
    {
        for (j = 0; j < N; j++)
        {

            if(fabs(A[i/NB][j/NB][(i%NB)*NB+j%NB]-Alin[i*N+j])> 1e-4)
            {
                cout<<"check inconsistent!"<<endl;
                cout<<"in tile ["<<i/NB<<"]["<<j/NB<<"], element "<<(i%NB)*NB+j%NB<<" is wrong."<<endl;
                cout<<"the value in Tiled A is : "<<A[i/NB][j/NB][(i%NB)*NB+j%NB]<<endl;
                cout<<"but in original ALin, it is "<<Alin[i*N+j]<<endl;
                return 1;
            }
        }
    }
    return 0;
}

void randn(int N, double* Alin)
{
    srand((float)time(0));
    int i,j;
    double temp;
    for(i=0; i<N*N; i++)
    {
        Alin[i]=rand()%2;
    }
    for(i=0; i<N; i++)
    {
        for(j=0; j<N; j++)
        {
            temp=Alin[i*N+j]+Alin[j*N+j];
            Alin[i*N+j]=temp;
            Alin[j*N+i]=temp;
        }
    }
    for(i=0; i<N; i++)
    {
        Alin[i*N+i]=0;
    }
}

void init(int NB, int DIM)
{
    int N = NB*DIM;
    // linear matrix
    Alin = (double *) malloc(N * N * sizeof(double));
    // fill the matrix with random values
    randn(N, Alin);

    int i;
    int j;
    // make it positive definite
    for(i=0; i<N; i++)
    {
        Alin[i*N + i] += N;
    }
    // blocked matrix
    A = (double ***) malloc(DIM*sizeof(double **));
    for (i = 0; i < DIM; i++)
    {
        A[i] = (double **) malloc(DIM*sizeof(double*));
        for(j=0; j<DIM; j++)
            A[i][j]=(double *)malloc(NB*NB*sizeof(double));
    }

    for (i = 0; i < N; i++)
    {
        for (j = 0; j < N; j++)
        {
            A[i/NB][j/NB][(i%NB)*NB+j%NB] = Alin[i*N+j];
        }
    }
}

int main(int argc, char *argv[])
{
    if (argc<3)
    {
        cout<<"please input DIM NB\n"<<endl;
        exit(0);
    }

    int DIM = atoi(argv[1]);
    int NB = atoi(argv[2]);
    //int num_threads = atoi(argv[3]);

    struct timeval start;
    struct timeval stop;
    unsigned long elapsed;
    double exctime;

    init(NB, DIM);
    //omp_set_num_threads(num_threads);

    gettimeofday(&start,NULL);
    //#pragma omp parallel num_threads(num_threads)
    //#pragma omp single
        compute(NB, DIM, A);
    
    gettimeofday(&stop,NULL);

    elapsed = 1000000*(stop.tv_sec - start.tv_sec);
    elapsed += stop.tv_usec - start.tv_usec;
    exctime=(double)elapsed/1000000;
    cout<< "Operations time: "<<exctime<<endl;

    char LO='L';
    int  INFO;
    int N = NB*DIM;
    //compute with mkl library
    dpotrf_(&LO, &N, Alin, &N, &INFO);
    //dpo(Alin, N);
    if(check(NB, N, Alin, A))
        exit(1);
    cout<<"check is passed, correct!"<<endl;

    return 0;
}
